import pandas as pd
import os
import sys
import re


"""

执行时间：
执行 
c-lop-analyze.py 
-> get_all_real_total_time_and_predict_time_to_summary_csv.py 
-> get_avg_precision_from_summary_csv.py 
之后进行排序输出

该脚本作用：
传入一个csv文件，根据csv文件中的proc进行排序
注意：只按进程数排序，不再考虑节点数

"""

def extract_number(text):
    """
    从字符串中提取第一个数字。
    
    Args:
        text (str): 包含数字的输入字符串
    
    Returns:
        int: 字符串中找到的第一个数字，如果没有找到数字则返回0
    """
    if pd.isna(text):
        return 0
    
    # 使用正则表达式在字符串中查找第一个数字
    match = re.search(r'\d+', str(text))
    if match:
        return int(match.group())
    return 0

def calculate_relative_error(df):
    """
    计算相对误差：|avg_total_real_comm_time - avg_predict_total_comm_time| / avg_total_real_comm_time * 100
    
    Args:
        df (pd.DataFrame): 包含数据的DataFrame
    
    Returns:
        pd.DataFrame: 添加了relative_error列的DataFrame
    """
    if 'avg_total_real_comm_time' in df.columns and 'avg_predict_total_comm_time' in df.columns:
        # 计算相对误差
        df['relative_error'] = (abs(df['avg_total_real_comm_time'] - df['avg_predict_total_comm_time']) / df['avg_total_real_comm_time']) * 100
        print("已添加relative_error列")
    else:
        print("警告: 缺少计算相对误差所需的列")
    
    return df

def sort_csv_by_proc(csv_path, output_path=None):
    """
    根据'proc_dir'列中的数值对CSV文件进行排序。
    只按进程数排序。
    同时计算相对误差。
    
    Args:
        csv_path (str): 输入CSV文件的路径
        output_path (str, optional): 保存排序后CSV文件的路径。如果为None，将覆盖输入文件。
    
    Returns:
        pd.DataFrame: 排序后的DataFrame
    """
    try:
        # 读取CSV文件
        df = pd.read_csv(csv_path)
        
        # 检查'proc_dir'列是否存在
        if 'proc_dir' not in df.columns:
            print(f"错误: CSV文件必须包含'proc_dir'列。")
            print(f"可用列: {df.columns.tolist()}")
            return None
        
        print(f"原始数据行数: {len(df)}")
        print(f"排序前前5行:")
        print(df[['proc_dir']].head())
        
        # 创建临时列，提取数值用于排序
        df['proc_num'] = df['proc_dir'].apply(extract_number)
        
        print(f"提取的数值:")
        print(df[['proc_dir', 'proc_num']].head())
        
        # 按数值对DataFrame进行排序（只按进程数排序）
        sorted_df = df.sort_values(by=['proc_num'])
        
        print(f"排序后前5行:")
        print(sorted_df[['proc_dir', 'proc_num']].head())
        
        # 删除临时列
        sorted_df = sorted_df.drop(['proc_num'], axis=1)
        
        # 计算相对误差
        sorted_df = calculate_relative_error(sorted_df)
        
        # 保存排序后的DataFrame
        if output_path:
            sorted_df.to_csv(output_path, index=False)
            print(f"排序后的CSV已保存到: {output_path}")
        else:
            sorted_df.to_csv(csv_path, index=False)
            print(f"原始CSV文件已被排序后的数据覆盖。")
        
        return sorted_df
    
    except Exception as e:
        print(f"处理CSV文件时出错: {e}")
        return None

if __name__ == "__main__":
   
    base_dir_csv = r"F:\PostGraduate\Point-to-Point-Code\App_Prediction\analysis_weak_scaling\analysis_for_100atom_per_proc\all_predict_precision.bak\avg_precision_summary.csv"

    # 检查文件是否存在
    if not os.path.exists(base_dir_csv):
        print(f"错误: 文件未找到: {base_dir_csv}")
        sys.exit(1)
    
    # 对CSV文件进行排序
    sort_csv_by_proc(base_dir_csv)
    print("CSV排序完成。")